//
// Created by YueChuan on 2023/4/13.
//

#include "socket.h"
#include "fd_manager.h"
#include "log.h"
#include "hook.h"
#include "macro.h"
#include <ctime>

static yuechuan::Logger::ptr g_logger = YUECHUAN_LOG_NAME("root");
namespace yuechuan {

    Socket::ptr Socket::CreateTCP(Address::ptr address) {
        Socket::ptr sock(new Socket(address->getFamily(), TCP, 0));
        return sock;
    }

    Socket::ptr Socket::CreateUDP(Address::ptr address) {
        Socket::ptr sock(new Socket(address->getFamily(), UDP, 0));
        sock->newSock();
        sock->m_isConnect = true;
        return sock;
    }

    Socket::ptr Socket::CreateTCPSocket() {
        Socket::ptr sock(new Socket(IPV4, TCP, 0));
        return sock;
    }

    Socket::ptr Socket::CreateTCPSocket6() {
        Socket::ptr sock(new Socket(IPV6, TCP, 0));
        return sock;
    }

    Socket::ptr Socket::CreateUDPSocket() {
        Socket::ptr sock(new Socket(IPV4, UDP, 0));
        sock->newSock();
        sock->m_isConnect = true;
        return sock;
    }

    Socket::ptr Socket::CreateUDPSocket6() {
        Socket::ptr sock(new Socket(IPV6, UDP, 0));
        sock->newSock();
        sock->m_isConnect = true;
        return sock;
    }

    Socket::ptr Socket::CreateUnixTCPSocket() {
        Socket::ptr sock(new Socket(UNIX, TCP, 0));
        return sock;
    }

    Socket::ptr Socket::CreateUnixUDPSocket() {
        Socket::ptr sock(new Socket(UNIX, UDP, 0));
        return sock;
    }

    Socket::Socket(int family, int type, int protocol) :
            m_sock(-1),
            m_family(family),
            m_type(type),
            m_protocol(protocol),
            m_isConnect(false) {

    }

    Socket::~Socket() {
        this->close();
    }

    bool Socket::init(int sock) {
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(sock);
        if (ctx && ctx->isSocket() && !ctx->isClose()) {
            m_sock = sock;
            m_isConnect = true;
            initSock();
            getLocalAddress();
            getRemoteAddress();
            return true;
        }
        return false;
    }


    int64_t Socket::getSendTimeout() const {
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(m_sock);
        if (ctx) {
            return ctx->getTimeout(SO_SNDTIMEO);
        }
        return -1;
    }

    void Socket::setSendTimeout(uint64_t timeout) {
        timeval tv = {time_t(timeout / 1000), suseconds_t(timeout % 1000 * 1000)};
        setOption(SOL_SOCKET, SO_SNDTIMEO, tv);
    }

    int64_t Socket::getRecvTimeout() const {
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(m_sock);
        if (ctx) {
            return ctx->getTimeout(SO_RCVTIMEO);
        }
        return -1;
    }

    void Socket::setRecvTimeout(uint64_t timeout) {
        timeval tv = {time_t(timeout / 1000), suseconds_t(timeout % 1000 * 1000)};
        setOption(SOL_SOCKET, SO_RCVTIMEO, tv);
    }

    bool Socket::getOption(int level, int optname, void *optval, size_t *optlen) {
        int rt = getsockopt(m_sock, level, optname, optval, (socklen_t *) optlen);
        if (rt) {
            YUECHUAN_LOG_ERROR(g_logger) << "getsockopt error sock = "
                                         << m_sock << "errno = " << errno << " " << strerror(errno);
            return false;
        }
        return true;
    }

    bool Socket::setOption(int level, int optname, const void *optval, socklen_t optlen) {
        int rt = setsockopt(m_sock, level, optname, optval, (socklen_t) optlen);
        if (rt) {
            YUECHUAN_LOG_ERROR(g_logger) << "setOption error sock = "
                                         << m_sock << " errno = " << errno << " " << strerror(errno)
                                         << " level = " << level << " option = " << optlen;
            return false;
        }
        return true;
    }

    Socket::ptr Socket::accept() {
        Socket::ptr sock(new Socket(m_family, m_type, m_protocol));
        int newSock = ::accept(m_sock, nullptr, nullptr);
        if (newSock == -1) {
            YUECHUAN_LOG_ERROR(g_logger) << "accept error sock = "
                                         << m_sock << "errno = " << errno << " " << strerror(errno);
            return nullptr;
        }
        if (sock->init(newSock)) {
            return sock;
        }
        return nullptr;
    }

    bool Socket::bind(const Address::ptr addr) {
        if (YUECHUAN_UNLIKELY(!isValid())) {
            newSock();
            if (YUECHUAN_UNLIKELY(!isValid())) {
                return false;
            }
        }
        if (addr->getFamily() != m_family) {
            YUECHUAN_LOG_ERROR(g_logger) << "bind socket.family = " << addr->getFamily()
                                         << " != m_family{" << m_family << "}"
                                         << " addr = " << addr->toString();
            return false;
        }
        int rt = ::bind(m_sock, addr->getAddress(), addr->getAddrLen());
        if (rt) {
            YUECHUAN_LOG_ERROR(g_logger) << "bind error m_sock = "
                                         << m_sock << " errno " << errno << " " << strerror(errno);
            return false;
        }
        getLocalAddress();
        return true;
    }

    bool Socket::connect(const Address::ptr addr, uint64_t timeout_ms) {
        m_remoteAddress = addr;
        if (!isValid()) {
            newSock();
            if (YUECHUAN_UNLIKELY(!isValid())) {
                return false;
            }
        }
        if (YUECHUAN_UNLIKELY(addr->getFamily() != m_family)) {
            YUECHUAN_LOG_ERROR(g_logger) << "connect socket.family " << addr->getFamily()
                                         << " != m_family{" << m_family << "}"
                                         << " addr = " << addr->toString();
            return false;
        }
        if (timeout_ms == (uint64_t) -1) {
            if (::connect(m_sock, addr->getAddress(), addr->getAddrLen())) {
                YUECHUAN_LOG_ERROR(g_logger) << "connect error m_sock = " << m_sock
                                             << " connect( " << addr->toString()
                                             << ") errno = " << errno << " " << strerror(errno);
                close();
                return false;
            }
        } else {
            if (connect_with_timeout(m_sock, addr->getAddress(), addr->getAddrLen(), timeout_ms)) {
                YUECHUAN_LOG_ERROR(g_logger) << "sock = " << m_sock << " connect(" << addr->toString()
                                             << " )timeout = " << timeout_ms << " error errno = " << errno
                                             << " " << strerror(errno);
                close();
                return false;
            }
        }
        m_isConnect = true;
        getRemoteAddress();
        getLocalAddress();
        return true;
    }

    bool Socket::reconnect(int timeout_ms) {
        if (!m_remoteAddress) {
            YUECHUAN_LOG_ERROR(g_logger) << " reconnect m_remoteAddress is null";
            return false;
        }
        m_localAddress.reset();
        return connect(m_remoteAddress, timeout_ms);
    }

    bool Socket::close() {
        if (!m_isConnect && m_sock == -1) {
            return true;
        }
        m_isConnect = false;
        if (m_sock != -1) {
            ::close(m_sock);
            m_sock = -1;
        }
        return true;
    }

    bool Socket::listen(int backlog) {
        if (!isValid()) {
            YUECHUAN_LOG_ERROR(g_logger) << "listen error sock = -1";
            return false;
        }
        if (::listen(m_sock, backlog)) {
            YUECHUAN_LOG_ERROR(g_logger) << " listen error  errno = "
                                         << errno << " " << strerror(errno);
            return false;
        }
        return true;
    }

    int Socket::send(const void *buffer, size_t length, int flags) {
        if (isConnect()) {
            return ::send(m_sock, buffer, length, flags);
        }
        return -1;
    }


    int Socket::send(const iovec *buffers, size_t length, int flag) {
        if (isConnect()) {
            msghdr msg;
            memset(&msg, 0, sizeof(msg));
            msg.msg_iov = (iovec *) buffers;
            msg.msg_iovlen = length;
            return ::sendmsg(m_sock, &msg, flag);
        }
        return -1;
    }

    int Socket::sendTo(const void *buffer, size_t length, Address::ptr to, int flags) {
        if (isConnect()) {
            return ::sendto(m_sock, buffer, length, flags, to->getAddress(), to->getAddrLen());
        }
        return -1;
    }

    int Socket::sendTo(const iovec *buffers, size_t length, Address::ptr to, int flags) {
        if (isConnect()) {
            msghdr msg;
            memset(&msg, 0, sizeof(msghdr));
            msg.msg_iov = (iovec *) buffers;
            msg.msg_iovlen = to->getAddrLen();
            msg.msg_name = to->getAddress();
            return ::sendmsg(m_sock, &msg, flags);

        }
        return -1;
    }

    int Socket::recv(void *buffer, size_t length, int flag) {
        if (isConnect()) {
            return ::recv(m_sock, buffer, length, flag);
        }
        return -1;
    }

    int Socket::recv(iovec *buffers, size_t length, int flag) {
        if (isConnect()) {
            msghdr msg;
            memset(&msg, 0, sizeof(msg));
            msg.msg_iov = (iovec *) buffers;
            msg.msg_iovlen = length;
            return ::recvmsg(m_sock, &msg, flag);
        }
        return -1;
    }

    int Socket::recvFrom(void *buffer, size_t length, Address::ptr from, int flag) {
        if (isConnect()) {
            socklen_t len = from->getAddrLen();
            return ::recvfrom(m_sock, buffer, length, flag, from->getAddress(), &len);
        }
        return -1;
    }

    int Socket::recvFrom(iovec *buffers, size_t length, Address::ptr from, int flag) {
        if (isConnect()) {
            msghdr msg;
            memset(&msg, 0, sizeof(msg));
            msg.msg_iovlen = length;
            msg.msg_iov = (iovec *) buffers;
            msg.msg_name = from->getAddress();
            msg.msg_iovlen = from->getAddrLen();
            return ::recvmsg(m_sock, &msg, flag);
        }
        return -1;
    }

    Address::ptr Socket::getRemoteAddress() {
        if (m_remoteAddress) {
            return m_remoteAddress;
        }
        Address::ptr result;
        switch (m_family) {
            case AF_INET:
                result.reset(new IPV4Address());
                break;
            case AF_INET6:
                result.reset(new IPV6Address());
                break;
            case AF_UNIX:
                result.reset(new UnixAddress());
                break;
            default:
                result.reset(new UnknowAddress(m_family));
                break;
        }
        socklen_t addrLen = result->getAddrLen();
        if (getpeername(m_sock, result->getAddress(), &addrLen)) {
            return Address::ptr(new UnknowAddress(m_family));
        }
        if (m_family == AF_UNIX) {
            UnixAddress::ptr addr = std::dynamic_pointer_cast<UnixAddress>(result);
            addr->setAddrLen(addrLen);
        }
        m_remoteAddress = result;
        return m_remoteAddress;

    }

    Address::ptr Socket::getLocalAddress() {
        if (m_localAddress) {
            return m_localAddress;
        }
        Address::ptr result;
        switch (m_family) {
            case AF_INET:
                result.reset(new IPV4Address());
                break;
            case AF_INET6:
                result.reset(new IPV6Address());
                break;
            case AF_UNIX:
                result.reset(new UnixAddress());
                break;
            default:
                result.reset(new UnknowAddress(m_family));
                break;
        }
        socklen_t addrLen = result->getAddrLen();
        if (getpeername(m_sock, result->getAddress(), &addrLen)) {
            return Address::ptr(new UnknowAddress(m_family));
        }
        if (m_family == AF_UNIX) {
            UnixAddress::ptr addr = std::dynamic_pointer_cast<UnixAddress>(result);
            addr->setAddrLen(addrLen);
        }
        m_localAddress = result;
        return m_localAddress;
    }

    bool Socket::isValid() const {
        return m_sock != -1;
    }

    int Socket::getError() {
        int error = 0;
        size_t len = sizeof(error);
        if (!getOption(SOL_SOCKET, SO_ERROR, &error, &len)) {
            error = errno;
        }
        return error;

    }

    std::ostream &Socket::dump(std::ostream &os) const {
        os << "[Socket sock = " << m_sock
           << " is_connected = " << m_isConnect
           << " family = " << m_family
           << " type = " << m_type
           << " protocol = " << m_protocol;
        if (m_localAddress) {
            os << " localAddress = " << m_localAddress->toString();
        }
        if (m_remoteAddress) {
            os << " remoteAdress = " << m_remoteAddress->toString();

        }
        os << "]";
        return os;
    }

    std::string Socket::toString() const {
        std::stringstream ss;
        dump(ss);
        return ss.str();
    }

    bool Socket::cancelRead() {
        return IOManager::GetThis()->cancelEvent(m_sock, IOManager::READ);
    }

    bool Socket::cancelWrite() {
        return IOManager::GetThis()->cancelEvent(m_sock, IOManager::WRITE);
    }

    bool Socket::cancelAccept() {
        return IOManager::GetThis()->cancelEvent(m_sock, IOManager::READ);
    }

    bool Socket::cancelAll() {
        return IOManager::GetThis()->cancelAll(m_sock);
    }

    void Socket::initSock() {
        int val = 1;
        setOption(SOL_SOCKET, SO_REUSEADDR, val);
        if (m_type == SOCK_STREAM) {
            setOption(IPPROTO_TCP, TCP_NODELAY, val);
        }
    }

    void Socket::newSock() {
        m_sock = socket(m_family, m_type, m_protocol);
        if (YUECHUAN_LIKELY(m_sock != -1)) {
            initSock();
        } else {
            YUECHUAN_LOG_ERROR(g_logger) << "new sock error "
                                         << errno << " " << strerror(errno);
        }
    }

    SSLSocket::SSLSocket(int family, int type, int protocol) : Socket(family, type, protocol) {

    }

    Socket::ptr SSLSocket::accept() {
        SSLSocket::ptr sock(new SSLSocket(m_family, m_type, m_protocol));
        int newsock = ::accept(m_sock, nullptr, nullptr);
        if (newsock == -1) {
            YUECHUAN_LOG_ERROR(g_logger) << "accept error = " << errno << " " << strerror(errno);
            return nullptr;
        }
        sock->m_ctx = m_ctx;
        if (sock->init(newsock)) {
            return sock;
        }
        return nullptr;
    }

    bool SSLSocket::bind(const Address::ptr addr) {
        return Socket::bind(addr);
    }

    bool SSLSocket::connect(const Address::ptr addr, uint64_t timeout_ms) {
        bool v = Socket::connect(addr, timeout_ms);
        if (v) {
            m_ctx.reset(SSL_CTX_new(SSLv23_client_method()), SSL_CTX_free);
            m_ssl.reset(SSL_new(m_ctx.get()), SSL_free);
            SSL_set_fd(m_ssl.get(), m_sock);
            v = (SSL_connect(m_ssl.get()) == 1);
        }
        return v;
    }

    bool SSLSocket::listen(int backlog) {
        return Socket::listen(backlog);
    }

    bool SSLSocket::close() {
        return Socket::close();
    }

    int SSLSocket::send(const void *buffer, size_t length, int flags) {
        if (m_ssl) {
            return SSL_write(m_ssl.get(), buffer, length);
        }
        return -1;
    }

    int SSLSocket::send(const iovec *buffers, size_t length, int flags) {
        if (!m_ssl) {
            return -1;
        }
        int total = 0;
        for (size_t i = 0; i < length; ++i) {
            int tmp = SSL_write(m_ssl.get(), buffers[i].iov_base, buffers[i].iov_len);
            if (tmp <= 0) {
                return tmp;
            }
            total += tmp;
            if (tmp != (int) buffers[i].iov_len) {
                break;
            }
        }
        return total;
    }

    int SSLSocket::sendTo(const void *buffer, size_t length, const Address::ptr to, int flag) {
        YUECHUAN_ASSERT(false);
        return -1;
    }

    int SSLSocket::sendTo(const iovec *buffers, size_t length, const Address::ptr to, int flags) {
        YUECHUAN_ASSERT(false);
        return -1;
    }

    int SSLSocket::recv(void *buffer, size_t length, int flag) {
        if (m_ssl) {
            return SSL_read(m_ssl.get(), buffer, length);
        }
        return -1;
    }

    int SSLSocket::recv(iovec *buffers, size_t length, int flags) {
        if (!m_ssl) {
            return -1;
        }
        int total = 0;
        for (size_t i = 0; i < length; ++i) {
            int tmp = SSL_read(m_ssl.get(), buffers[i].iov_base, buffers[i].iov_len);
            if (tmp <= 0) {
                return tmp;
            }
            total += tmp;
            if (tmp != (int) buffers[i].iov_len) {
                break;
            }
        }
        return total;
    }

    int SSLSocket::recvFrom(void *buffer, size_t length, Address::ptr addr, int flag) {
        YUECHUAN_ASSERT(false);
        return -1;
    }

    int SSLSocket::recvFrom(iovec *buffers, size_t length, Address::ptr addr, int flags) {
        YUECHUAN_ASSERT(false);
        return -1;
    }

    bool SSLSocket::loadCertificates(const std::string &cert_file, const std::string &key_file) {
        m_ctx.reset(SSL_CTX_new(SSLv23_server_method()), SSL_CTX_free);
        if (SSL_CTX_use_certificate_chain_file(m_ctx.get(), cert_file.c_str()) != 1) {
            YUECHUAN_LOG_ERROR(g_logger) << "SSL_CTX_use_certificate_chain_file ("
                                         << cert_file << "error ";
            return false;
        }
        if (SSL_CTX_use_PrivateKey_file(m_ctx.get(), key_file.c_str(), SSL_FILETYPE_PEM) != 1) {
            YUECHUAN_LOG_ERROR(g_logger) << "SSL_CTX_use_PrivateKey_file("
                                         << key_file << "error ";
            return false;
        }
        if (SSL_CTX_check_private_key(m_ctx.get()) != 1) {
            YUECHUAN_LOG_ERROR(g_logger) << "SSL_CTX_check_private_key cert_file ("
                                         << cert_file << "error ";
            return false;
        }
        return true;
    }

    SSLSocket::ptr SSLSocket::CreateTcp(Address::ptr address) {
        SSLSocket::ptr sock(new SSLSocket(address->getFamily(), TCP, 0));
        return sock;
    }

    SSLSocket::ptr SSLSocket::CreateTcpSocket() {
        SSLSocket::ptr sock(new SSLSocket(IPV4, TCP, 0));
        return sock;
    }

    SSLSocket::ptr SSLSocket::CreateTcpSocket6() {
        SSLSocket::ptr sock(new SSLSocket(IPV6, TCP, 0));
        return sock;
    }

    bool SSLSocket::init(int sock) {
        bool v = Socket::init(sock);
        if (v) {
            m_ssl.reset(SSL_new(m_ctx.get()), SSL_free);
            SSL_set_fd(m_ssl.get(), m_sock);
            v = (SSL_accept(m_ssl.get()) == 1);
        }
        return v;
    }

    std::ostream &SSLSocket::dump(std::ostream &os) const {
        os << "[SSLSocket sock = " << m_sock
           << "is_connect = " << m_isConnect
           << " family = " << m_family
           << " type = " << m_type
           << " protocol" << m_protocol;
        if (m_localAddress) {
            os << " localAddress = " << m_localAddress->toString();
        }
        os << " ]";
        return os;
    }

    std::ostream &operator<<(std::ostream &os, const Socket &sock) {
        return sock.dump(os);
    }


}// namespace yuechuan
